import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;
/**
 * implements a data structure used for encoding and decoding Strings with Huffman algorithm.
 * @author baotam
 *
 */
public class HuffmanTree {

	private Node root;
	private Map<Character, String> encodingTable;
	
	/**
	 * generate a Huffman tree with the given frequency table
	 * @param frequencies
	 */
	public HuffmanTree(Map<Character, Integer> frequencies) {
		if(frequencies.size() == 0)	{
			encodingTable = new HashMap<Character,String>();
			return;
		}
		//pre-process the nodes with the frequency table
		PriorityQueue<Node> nodeQueue = new PriorityQueue<Node>();
		for(Character currentChar : frequencies.keySet()) {
			nodeQueue.add(new Node(currentChar, frequencies.get(currentChar)));
		}
		//generate the tree
		while(nodeQueue.size() > 1) {
			nodeQueue.add(new Node(nodeQueue.poll(), nodeQueue.poll()));
		}
		//add in the root
		root = nodeQueue.poll();
		//generate encoding table
		encodingTable = generateEncodingTable();
		
	}
	
	/**
	 * 
	 * @return a table of characters and their encoded values
	 */
	public Map<Character, String> generateEncodingTable() {
		
		/*
		 * You will need to write a method to traverse the tree to generate this
		 * encoding table. Think pre-order traversal, but you add to a map instead
		 * of a list
		 */
		Map<Character, String> encodingTable = new HashMap<Character, String>();
		generateEncodingTable(encodingTable, root, "");
		return encodingTable;
	}
	
	/**
	 * preorderedly traverse the Huffman tree and populate the encoding table
	 * @param encodingTable
	 * @param node
	 * @param code
	 */
	private void generateEncodingTable(Map<Character, String> encodingTable, Node node, String code) {
		if(node.getData() != null) {
			encodingTable.put(node.getData(), code);
			return;
		}
		generateEncodingTable(encodingTable, node.getLeft(), code + "0");
		generateEncodingTable(encodingTable, node.getRight(), code + "1");
	}
	
	/**
	 * encode the given String
	 * @param input
	 * @return the encoded String of '1' and '0'
	 */
	public String encode(String input) {
		String result = "";
		for(int i = 0; i < input.length(); i++) {
			result += encodingTable.get(input.charAt(i));
		}
		return result;
	}
	
	/**
	 * decode the given String
	 * @param output
	 * @return the decoded String
	 */
	public String decode(String output) {
		String result = "";
		int currentIndex = 0;
		while(currentIndex < output.length()) {
			Character decodedChar = decode(root, output, currentIndex);
			result += decodedChar;
			currentIndex += encodingTable.get(decodedChar).length();
		}
		return result;
	}
	
	/**
	 * Search the tree for the character at the given index of the "output" String
	 * @param node
	 * @param output
	 * @param currentIndex
	 * @return the character at "currentIndex"
	 */
	private Character decode(Node node, String output, int currentIndex) {
		if(node.getData() != null) {
			return node.getData();
		}
		
		if(output.charAt(currentIndex) == '0') {
			return decode(node.getLeft(), output, currentIndex+1);
		}
		else {
			return decode(node.getRight(), output, currentIndex+1);
		}
	}
	
	
	/*
	 * Do not modify the methods below this point, they are used for grading.
	 */

	public void setRoot(Node data) {
		this.root = data;
	}

	public Node getRoot() {
		return root;
	}
}
